# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
import random  # 随机数生成模块，用于设置随机种子和生成随机值
from typing import List, Optional, Literal  # 类型标注模块，提供列表、可选值、字面值类型支持
from enum import Enum, EnumMeta  # 枚举类型模块，用于创建枚举类和元类
from dataclasses import dataclass, field  # 数据类模块，用于简化类定义和字段配置


class StrEnumMeta(EnumMeta):
    """字符串枚举元类 - 解决 Hydra 配置框架中的实例检查问题"""
    # 这是针对 submitit 序列化导致 hydra 中 StrEnum 实例检查失败的解决方案
    # 参考: https://github.com/facebookresearch/hydra/issues/1156
    @classmethod
    def __instancecheck__(cls, other):
        """重写实例检查方法，通过检查类型字符串中是否包含'enum'来判断"""
        return "enum" in str(type(other))


class StrEnum(Enum, metaclass=StrEnumMeta):
    """字符串枚举基类 - 提供字符串形式的枚举值比较和表示"""
    def __str__(self):
        """返回枚举值的字符串表示"""
        return self.value

    def __eq__(self, other: str):
        """支持与字符串直接比较相等性"""
        return self.value == other

    def __repr__(self):
        """返回枚举值的表示形式"""
        return self.value

    def __hash__(self):
        """返回枚举值的哈希值，用于字典键等场景"""
        return hash(str(self))


def ChoiceEnum(choices: List[str]):
    """动态创建字符串选择枚举类 - 根据字符串列表生成对应的枚举类型"""
    return StrEnum("Choices", {k: k for k in choices})


_SEARCH_CHOICES = [
    "mcts",         # 蒙特卡洛树搜索 (Monte Carlo Tree Search)
    "mcts_wo_code", # 无代码模式的蒙特卡洛树搜索 (MCTS without code execution)
    "mcts_caption", # Caption 扩展的蒙特卡洛树搜索 (MCTS with Caption actions)
    "bs",           # 束搜索 (Beam Search)
]

SEARCH_CHOICES = ChoiceEnum(_SEARCH_CHOICES)  # 搜索算法选择枚举类


_PROMPT_CHOICES = [
    "rstar",  # rStar 提示词格式
]

PROMPT_CHOICES = ChoiceEnum(_PROMPT_CHOICES)  # 提示词格式选择枚举类
@dataclass
class BaseConfig:
    """
rStar-Math 推理系统的基础配置类
包含搜索算法、模型路径、推理参数、MCTS参数等所有核心配置项
"""

    # === 核心模式和路径配置 ===
    mode: SEARCH_CHOICES = field(
        default="mcts", metadata={"help": "推理搜索模式选择"}
    )  # 选择使用的搜索算法：mcts(蒙特卡洛树搜索) 或 bs(束搜索)
    model_dir: Optional[str] = field(
        default=None, metadata={"help": "语言模型目录路径"}
    )  # 策略模型(LLM)的本地路径或HuggingFace模型名称
    reward_model_dir: Optional[str] = field(
        default=None, metadata={"help": "奖励模型目录路径"}
    )  # PPM过程偏好模型的路径，用于评估推理步骤质量
    few_shot_path: Optional[str] = field(
        default=None, metadata={"help": "少样本示例数据文件路径"}
    )  # 包含few-shot示例的JSON文件路径
    prompt_path: Optional[str] = field(
        default=None, metadata={"help": "提示词配置文件路径"}
    )  # 自定义提示词模板的JSON配置文件路径
    num_few_shot: int = field(
        default=0, metadata={"help": "使用的少样本示例数量"}
    )  # 在提示词中包含的few-shot示例个数
    
    # === 新增：模型后端配置 ===
    backend_type: str = field(
        default="vllm", metadata={"help": "模型后端类型: vllm, openai"}
    )  # 模型调用后端类型，支持本地vLLM或远程OpenAI API
    models_config_path: str = field(
        default="config/models_config.yaml", metadata={"help": "模型配置文件路径"}
    )  # 模型配置文件路径，包含API密钥等敏感信息
    
    # === 新增：多模态配置 ===
    multimodal_enabled: bool = field(
        default=False, metadata={"help": "是否启用多模态功能"}
    )  # 是否启用图像+文本的多模态推理能力
    image_root: str = field(
        default="./images", metadata={"help": "图像文件根目录"}
    )  # 图像文件存储的根目录路径
    max_image_size: List[int] = field(
        default_factory=lambda: [1024, 1024], metadata={"help": "最大图像尺寸[width, height]"}
    )  # 图像预处理时的最大尺寸限制
    
    # === 新增：数据字段映射配置 ===
    question_field: str = field(
        default="question", metadata={"help": "问题字段名"}
    )  # 数据集中问题字段的名称
    answer_field: str = field(
        default="answer", metadata={"help": "答案字段名"}
    )  # 数据集中答案字段的名称
    image_field: str = field(
        default="image", metadata={"help": "图像字段名"}
    )  # 数据集中图像字段的名称
    
    # === 提示词格式配置 ===
    prompt_wrap: PROMPT_CHOICES = field(
        default="rstar", metadata={"help": "提示词包装格式类型"}
    )  # 提示词的格式化方式，目前支持rstar格式
    result_unwrap: PROMPT_CHOICES = field(
        default="rstar", metadata={"help": "结果解析格式类型"}
    )  # 模型输出结果的解析方式
    step_delim: str = field(
        default="\n", metadata={"help": "推理步骤间的分隔符"}
    )  # 用于分隔不同推理步骤的字符串
    
    # === vLLM推理引擎参数 ===
    temperature: float = field(
        default=0.7, metadata={"help": "控制生成文本的随机性"}
    )  # 温度参数，越高生成越随机，越低越确定性
    top_p: float = field(
        default=1.0, metadata={"help": "核采样参数，控制累积概率阈值"}
    )  # Nucleus采样的概率阈值，必须在(0,1]范围内
    top_k: int = field(
        default=-1, metadata={"help": "Top-K采样参数，控制候选token数量"}
    )  # 每步采样时考虑的最高概率token数量，-1表示不限制
    use_beam_search: bool = field(
        default=False, metadata={"help": "是否启用束搜索解码"}
    )  # 是否在vLLM层面使用束搜索（注意与mode='bs'的区别）
    best_of: int = field(
        default=1, metadata={"help": "束搜索中考虑的候选序列数量"}
    )  # 束搜索解码过程中的候选序列数
    max_tokens: int = field(
        default=2048, metadata={"help": "每个输出序列的最大token数"}
    )  # 单次生成的最大token长度限制
    seed: Optional[int] = field(
        default=None, metadata={"help": "随机种子，用于可重现的生成"}
    )  # 设置随机种子以确保结果可重现，None表示随机
    swap_space: Optional[int] = field(
        default=8, metadata={"help": "vLLM的交换空间大小(GB)"}
    )  # vLLM引擎使用的GPU-CPU交换空间
    
    # === 搜索算法核心参数 ===
    # 注意：n_generate_sample 在 Caption 模式下会被动态计算为 reason_samples + caption_samples
    n_generate_sample: int = field(
        default=4, metadata={"help": "每步生成的候选样本数量(论文中的B2)"}
    )  # 在每个搜索节点扩展时生成的候选步骤数
    
    # === Caption 扩展配置 ===
    reason_samples: int = field(
        default=3, metadata={"help": "Reason 动作生成的节点数量"}
    )  # Reason 动作在每次扩展时生成的候选步骤数
    caption_samples: int = field(
        default=1, metadata={"help": "Caption 动作生成的节点数量"}
    )  # Caption 动作在每次扩展时生成的候选步骤数
    
    stop: Optional[List[str]] = field(
        default=None, metadata={"help": "推理步骤的可能停止标记列表"}
    )  # 模型生成时的停止token列表
    step_beam_width: int = field(
        default=1, metadata={"help": "每步的束宽度(论文中的B1)"}
    )  # 束搜索模式下每步保留的最佳路径数
    max_depth: int = field(
        default=4, metadata={"help": "搜索树的最大深度(最大推理步数)"}
    )  # 限制推理过程的最大步骤数，防止无限循环
    iterations: int = field(
        default=1, metadata={"help": "MCTS的模拟迭代次数"}
    )  # MCTS算法中的rollout次数
    
    # === 奖励和评估参数 ===
    positive_reward: float = field(
        default=1.0, metadata={"help": "正确答案的奖励值"}
    )  # 推理成功时给予节点的奖励分数
    negative_reward: float = field(
       default=-1.0, metadata={"help": "错误答案的惩罚值"}
    )  # 推理失败时给予节点的惩罚分数
    errors_threshold: int = field(
        default=0, metadata={"help": "允许的最大连续代码错误数"}
    )  # 超过此错误数时终止该节点的扩展
    need_value_func: bool = field(
        default=False, metadata={"help": "是否在解码过程中使用价值函数"}
    )  # 是否启用PPM奖励模型进行步骤价值估计
    update_leaf_value: bool = field(
        default=False, metadata={"help": "是否在MCTS中更新叶节点价值"}
    )  # 是否对新扩展的叶节点进行价值更新
    
    # === MCTS特定参数 ===
    c_puct: float = field(
        default=2, metadata={"help": "MCTS中PUCT公式的探索常数"}
    )  # 控制MCTS选择过程中探索与利用的平衡
    is_sampling: bool = field(
        default=False, metadata={"help": "MCTS中是否使用采样模式生成解"}
    )  # 是否在解生成过程中使用随机采样
    prune: bool = field(
        default=False, metadata={"help": "是否在完整MCTS树中进行剪枝"}
    )  # 离线推理时是否剪枝搜索树以节省内存
    terminal_sample: bool = field(
        default=False
    )  # 是否对终端节点进行采样（实验性功能）
    
    # === 系统资源和批处理参数 ===
    batch_size: int = field(
        default=-1, metadata={"help": "批量推理的批次大小"}
    )  # 批量处理时的样本数，-1表示处理所有数据
    max_model_len: int = field(
        default=8192, metadata={"help": "模型的最大序列长度"}
    )  # 模型能处理的最大输入长度
    llm_gpu_memory_utilization: float = field(
        default=0.5, metadata={"help": "策略模型的GPU内存利用率"}
    )  # GPU内存分配比例，需要为奖励模型预留约15GB内存
    # 建议设置：80GB显存设为0.7，40GB显存设为0.5
    tp: int = field(
        default=1, metadata={"help": "语言模型的张量并行大小"}
    )  # 模型推理时的并行度，用于加速大模型推理
    
    # === 调试和输出参数 ===
    save_intermediate_rollouts: bool = field(
        default=True, metadata={"help": "是否保存中间rollout结果到./rollout目录"}
    )  # 是否保存每次rollout的中间结果用于分析调试
    debug_prompts: bool = field(
        default=False, metadata={"help": "是否输出prompt调试信息到logs目录"}
    )  # 是否保存发送给LLM的prompt和返回的结果，用于调试
    
    # === 新增：数据集处理配置 ===
    dataset_start_idx: int = field(
        default=0, metadata={"help": "数据集处理起始索引"}
    )  # 从数据集的第几条数据开始处理，用于断点续传或分片处理
    dataset_max_samples: int = field(
        default=-1, metadata={"help": "最大处理样本数，-1表示处理全部"}
    )  # 最多处理多少条数据，-1表示处理全部数据
    dataset_shuffle: bool = field(
        default=False, metadata={"help": "是否随机打乱数据集顺序"}
    )  # 是否在处理前随机打乱数据集
    dataset_seed: Optional[int] = field(
        default=None, metadata={"help": "数据集随机种子，用于可重现的打乱"}
    )  # 数据集打乱的随机种子，确保结果可重现
    
    # === 新增：vLLM 引擎增强配置 ===
    vllm_max_requests: int = field(
        default=1000, metadata={"help": "vLLM引擎最大并发请求数"}
    )  # vLLM引擎能同时处理的最大请求数量
    vllm_max_seq_len_to_capture: int = field(
        default=8192, metadata={"help": "vLLM引擎捕获的最大序列长度"}
    )  # vLLM引擎用于优化的最大序列长度
    vllm_disable_log_stats: bool = field(
        default=True, metadata={"help": "是否禁用vLLM统计日志"}
    )  # 是否禁用vLLM的详细统计日志输出
    vllm_enable_chunked_prefill: bool = field(
        default=True, metadata={"help": "是否启用vLLM分块预填充"}
    )  # 是否启用分块预填充以提高内存效率
    
    # === 新增：4×A100 优化参数 ===
    vllm_enable_prefix_caching: bool = field(
        default=False, metadata={"help": "是否启用vLLM前缀缓存"}
    )  # 启用前缀缓存，对MCTS重复前缀有巨大性能提升
    vllm_max_num_seqs: int = field(
        default=256, metadata={"help": "vLLM引擎最大序列数"}
    )  # 同时处理的最大序列数量，影响并发性能
    vllm_max_num_batched_tokens: int = field(
        default=8192, metadata={"help": "vLLM引擎批处理token数"}
    )  # 单批次处理的最大token数量，影响吞吐量
    rollout_parallel_workers: int = field(
        default=1, metadata={"help": "rollout并行工作进程数"}
    )  # 并行处理rollout的工作进程数量，1表示串行处理
    step_parallel_workers: int = field(
        default=12, metadata={"help": "step后处理并行工作进程数"}
    )  # step后处理阶段的并行工作进程数量